'''
Created on 2011-8-29

@author: Administrator
'''
import datetime
import protocol
import task
import logging
import trigger
import tasktemplate
import db
import servertime
import pickle
import os
import xml.etree.ElementTree 

PICKLEFILE = 'runingtaskpickle.pl'

WAITING_TIME_INTERVAL = datetime.timedelta(seconds=10)

class Scheduler():
    index = 1
    
    def __init__(self, servercfg, clientmanager, connectmanager, dbmanager, resultprocessor):
        self.tasks = {}
        self.parser_processes = []
        self.task_templates = {}
        self.pending_tasks = {}
        self.waitingruntemplate = {}
        self.task_finish_handler = []
        
        self.clientmanager = clientmanager
        self.connectmanager = connectmanager
        self.dbmanager = dbmanager
        self.resultprocessor = resultprocessor
        self.servercfg = servercfg
        
        self.HandData()
        
        
    def AddTaskFinishHander(self, handler):
        self.task_finish_handler.append(handler)
        
    
    def HandData(self):
        self.connectmanager.RegisterProtocol(protocol.c2s_task_start_respond, self.OnStartTaskRespond)
        self.connectmanager.RegisterProtocol(protocol.c2s_task_finish, self.OnTaskFinish)
        self.connectmanager.RegisterProtocol(protocol.c2s_terminated_task_respond, self.OnTerminatedTaskRespond)
        self.connectmanager.RegisterProtocol(protocol.c2s_running_task, self.OnRuningTaskInfoRequest)
        
        self.dbmanager.RegisterProtocol(db.dbprotocol.get_case_respond, self.GetCaseRespond)
        self.dbmanager.RegisterProtocol(db.dbprotocol.insert_task_respond, self.InsertTaskRespond)
        self.dbmanager.RegisterProtocol(db.dbprotocol.get_tasktemplate_respond, self.OnUpdateTaskTemplateRespond)
        
        self.clientmanager.disconnect_handler.append(self.HandleDisconnect)
     
        
    def HandleDisconnect(self, machineid, tasklist):
        for k, v in tasklist.items():
            if k not in self.tasks:
                logging.warn('taskid not exits')
                continue
        
            t = self.tasks[k]
            
            for subid in v:
                if subid not in t.subtasks:
                    logging.warn('subtaskid not exits')
                    continue
        
                subtask = t.subtasks[subid]
                
                if subtask.run_machine != machineid:
                    logging.warn('machine id not match')
                    continue
                
                subtask.start_disconnect_time = servertime.current_time
    
    def AddTask(self, task):
        self.tasks[self.index] = task
        self.index = self.index + 1
    
    
    def Uninit(self):
        f = None
        try:
            if len(self.tasks) != 0:
                f = open(PICKLEFILE, 'wb')
                for v in self.tasks.values():
                    v.start_disconnect_time = servertime.current_time
                pickle.dump(self.tasks, f)
        except Exception :
            logging.exception('dump tasks fail')
        finally:
            if f != None:
                f.close()
        
    def Init(self):
        
        all_task_template = self.dbmanager.LoadAllCaseTask()
        now = datetime.datetime.now()
        for v in all_task_template:
            self.LoadTaskTemplate(v)
        
        f = None
        try:
            if os.path.isfile(PICKLEFILE):
                f = open(PICKLEFILE, 'rb')
                tasks = pickle.load(f)
                for v in tasks.values():
                    if v.template.id not in self.task_templates:
                        logging.warn('task templateid not exits')
                        continue
                    if v.start_disconnect_time == None:
                        logging.warn('task start disconnect time must not None')
                        continue
                    if  (now - v.start_disconnect_time) > datetime.timedelta(hours=2):
                        logging.warn('task is timeouts')
                        continue
                    v.id = self.index
                    self.tasks[v.id] = v
                    self.index += 1
        except Exception :
            logging.exception('load last running tasks fail') 
        finally:
            if f != None:
                f.close()
                os.remove(PICKLEFILE)
               
                
    def DealTriggerTime(self, info):
        timetrigger = None
        try:
            if info.ExecuteTimeCron != None:
                cronexpression = info.ExecuteTimeCron
                start_time = info.BeginRunTime
                vlues = trigger.CronTrigger.CronExpression2Args(cronexpression)
                vlues['start_date'] = start_time
                timetrigger = trigger.CronTrigger(**vlues)
                
            if info.NextRunTime != None and timetrigger == None:
                return None
            
            if timetrigger == None:
                timetrigger = trigger.SimpleTrigger(info.BeginRunTime)
        except Exception as msg:
            logging.exception(msg)     
        return timetrigger   
        
        
    def LoadTaskTemplate(self, info):
        try:
            timetrigger = None
            if info.ExecuteTimeCron != None:
                cronexpression = info.ExecuteTimeCron
                start_time = info.BeginRunTime
                vlues = trigger.CronTrigger.CronExpression2Args(cronexpression)
                vlues['start_date'] = start_time
                timetrigger = trigger.CronTrigger(**vlues)
                
            if info.NextRunTime != None and timetrigger == None:
                return 
            
            if timetrigger == None:
                timetrigger = trigger.SimpleTrigger(info.BeginRunTime)
            

            machineidlist = []
            if info.MachineIds != None:
                for machine_id in info.MachineIds.split(','):
                    try:
                        machineid = int(machine_id)
                        machineidlist.append(machineid)
                    except ValueError:
                        logging.exception('conver machines error')
            machinecondition = tasktemplate.MachineCondition(machineidlist, info.OSId, info.MachineGroupId)
            
            script_path = info.ScriptPath
            sys_args = info.Argument
            
            if info.ScriptPath and info.ScriptPath == 'CustomScript':
                try:
                    tree = xml.etree.ElementTree.fromstring(info.Argument.encode('utf-8'))
                    script_path = tree.getiterator('CustomScript')[0].get('Path')
                    sys_args = tree.getiterator('CustomScript')[0].get('Argument')
                except Exception:
                    logging.exception('parse custom script error')
                    return
            
            if info.CaseCondition != None and len(info.CaseCondition) != 0:
                template = tasktemplate.CaseTaskTemplate(info.Id, machinecondition,
                                                         info.Name, script_path,
                                                         timetrigger, info.CaseCondition,
                                                         info.EmailReceiver,
                                                         sys_args,
                                                         next_run_time=info.NextRunTime,
                                                         scxml=info.SCXML,
                                                         success_receiver=info.SuccessMailReceiver,
                                                         fail_receiver=info.FailMailReceiver,
                                                         except_receiver=info.ExceptMailReceiver)
            else:
                template = tasktemplate.TaskTemplate(info.Id, machinecondition,
                                                     info.Name, script_path,
                                                     timetrigger,
                                                     info.EmailReceiver,
                                                     sys_args,
                                                     next_run_time=info.NextRunTime,
                                                     scxml=info.SCXML,
                                                     success_receiver=info.SuccessMailReceiver,
                                                     fail_receiver=info.FailMailReceiver,
                                                     except_receiver=info.ExceptMailReceiver)
        
            nextruntime = template.GetNextRunTime()
            if nextruntime == None:
                template.SetNextRunTime(servertime.current_time)
                
            if template.id in self.task_templates:
                logging.info('update tasktemplate %d' % template.id)
            else:
                logging.info('new tasktemplate  %d' % template.id)
                
            self.task_templates[template.id] = template
        except Exception:
            logging.exception('load task template failure')
            return False
        else:
            return True
        
    def RemoveTask(self, index):
        if index not in self.tasks:
            raise ValueError('no exits job id %d' % index)
        
        self.tasks.pop(index)

        
    def Loop(self):
        now = servertime.current_time
        for v in self.task_templates.values():
            next_run_time = v.GetNextRunTime()
            if next_run_time == None:
                continue
            if v.GetNextRunTime() < now:
                self.index = self.index + 1
                if self.NewTask(v.id,):
                    v.SetNextRunTime(now)
        
        finishtask = []
        for v in self.tasks.values():
            if v.id == None:
                continue
            if v.GetStatus() == task.IDLE:
                v.GettingStart(self)
            elif v.GetStatus() == task.WAITINGTORUN and v.AllReadyToRun():
                if now - v.create_time >= WAITING_TIME_INTERVAL:
                    v.SetTimeout(now)
                    continue
                machinelist = self.clientmanager.GetClientByCondition(v.template.machine_info)
                if len(machinelist) != 0:
                    taskinfos = v.Start(machinelist, now)
                    for taskinfo in taskinfos:
                        logging.info(taskinfo.args)
                        self.DoStartTask(taskinfo.machine.connectid, v.template.id, v.id, \
                                         taskinfo.subid, taskinfo.type, taskinfo.info, taskinfo.args)
                    self.UpdateTask(v)
                    self.UpdateSubTask(v.subtasks.values())
            elif v.GetStatus() == task.RUNING:
                v.Breath()
            elif v.IsFinish():
                finishtask.append(v)
            
        for v in finishtask:
            self.resultprocessor.PutTaskResult(v)
            self.UpdateTask(v)
            logging.info('template %s remove task %s' % (v.template.id, v.id))
            
            for handler in self.task_finish_handler:
                handler(v)
                
            v.template.RemoveTask(v.tempid)
            self.tasks.pop(v.id)
        
    
    def DoStartTask(self, clientid, templateid, taskid, subid, type, info, args):
        logging.info("start task %s:%s on %s" % (taskid, subid, clientid))
        request = protocol.S2CStartTask()
        request.templateid = templateid
        request.taskid = taskid
        request.subid = subid
        request.type = type
        request.info = info
        request.args = args
        logging.info(args)
        self.connectmanager.DoSend(clientid, request)

    
    def OnStartTaskRespond(self, clientid, respond):
        try:
            machine = self.clientmanager.GetClient(clientid)
        except ValueError as e:
            logging.warn('{0} not login client'.format(clientid))
            return
        
        if respond.taskid not in self.tasks:
            logging.warn('taskid not exits tasks')
            return
        
        if respond.result == False:
            t = self.tasks[respond.taskid]
            t.Finish(respond.subid, datetime.datetime.now(), task.EXCEPTION, respond.message)
            machine = self.clientmanager.GetClient(clientid)
            machine.RemoveTask(t.template.id, respond.taskid, respond.subid)
            
    
    def OnTaskFinish(self, clientid, respond):
        try:
            machine = self.clientmanager.GetClient(clientid)
        except ValueError as e:
            logging.warn('{0} not login client'.format(clientid))
            return
        
        if respond.taskid not in self.tasks:
            logging.warn('not exits taskid')
            return
        
        logging.info("finish task %s:%s" % (respond.taskid, respond.subid))
        finishtask = self.tasks.get(respond.taskid)
        
        machine = self.clientmanager.GetClient(clientid)
        machine.RemoveTask(finishtask.template.id, respond.taskid, respond.subid)
        
        finishtask.Finish(respond.subid, datetime.datetime.now(), respond) 
        
        subtask = finishtask.subtasks[respond.subid]
        self.UpdateSubTask([subtask])
        
        
    def CaseDataRequset(self, taskid, casesql):
        request = db.dbprotocol.GetCaseRequest()
        request.taskid = taskid
        request.sql = casesql
        self.dbmanager.PutRequest(request)
        
    def GetCaseRespond(self, respond):
        if respond.taskid not in self.tasks:
            logging.warn('{0} not exits'.format(respond.taskid))
            return 
        t = self.tasks[respond.taskid]
        if t.GetStatus() != task.WAITINGTORUN:
            logging.warn('error task status')
            return
        print(respond.cases)
        t.SetCase(respond.cases)
        
    def InsertTaskInDB(self, new_task):
        request = db.dbprotocol.InsertTaskRequest()
        request.taskid = new_task.tempid
        request.templateid = new_task.template.id
        request.status = 1
        self.dbmanager.PutRequest(request)
    
    def InsertTaskRespond(self, respond):
        if respond.taskid not in self.pending_tasks:
            logging.warn('error task temp id ')
            return
        
        t = self.pending_tasks.pop(respond.taskid)
        if respond.result == False:
            logging.warn('task insert db error')
            return 
        t.id = respond.dbid
        self.tasks[t.id] = t
        
    def GetTaskTemplate(self, id):
        if id in self.task_templates:
            return self.task_templates[id]
    
    
    def UpdateTask(self, task):
        request = db.dbprotocol.UpdateTaskRequest()
        request.taskid = task.id
        request.data = {'StartTime':task.begin_run_time,
                         'EndTime':task.end_time,
                         'ResultStatusId': task.result,
                         'TaskStatusId':task.GetStatus(),
                         'ResultMessage':task.result_message}
        self.dbmanager.PutRequest(request)

    def NewTask(self, templateid, casesql=None, maillist=None):
        if templateid not in self.task_templates:
            raise ValueError('id not in task')
        
        self.index = self.index + 1
        template = self.task_templates[templateid]
        new_task = template.CreateTask(self.index)
        if new_task == None:
            return False, None
        
        new_task.create_time = datetime.datetime.now() # Add
        if new_task.tempid in self.pending_tasks:
            logging.warn('same id in list')
            return False, None
        
        self.pending_tasks[new_task.tempid] = new_task
        
        if casesql != None:
            new_task.casesql = casesql
            
        self.InsertTaskInDB(new_task)
        
        return True, new_task
    
    
    def DoTerminatedTask(self, connectid, taskid, subid):
        request = protocol.S2CTerminatedTask(taskid, subid)
        self.connectmanager.DoSend(connectid, request)
        
        
    def OnTerminatedTaskRespond(self, clientid, respond):
        try:
            machine = self.clientmanager.GetClient(clientid)
        except ValueError as e:
            logging.warn('{0} not login client'.format(clientid))
            return
        
        if respond.taskid not in self.tasks:
            logging.warn('not exits task id')
            return
        
        t = self.tasks[respond.taskid]
        
        if t.GetStatus() != task.CANCELING:
            logging.warn('error task status')
            return
        
        if respond.result == False:
            logging.warn('terminated task fail %s' % respond.message)
            return
        
        logging.info("Terminated task %s:%s" % (respond.taskid, respond.subid))
        finishtask = self.tasks.get(respond.taskid)
        
        machine = self.clientmanager.GetClient(clientid)
        machine.RemoveTask(finishtask.template.id, respond.taskid, respond.subid)
        
        finishtask.OnSubTerminatedInfo(respond.subid, datetime.datetime.now()) 
        
        subtask = finishtask.subtasks[respond.subid]
        self.UpdateSubTask([subtask])
            
    def UpdateTaskTemplate(self, id):
        request = db.dbprotocol.GetTaskTemplateInfo()
        request.templateid = id
        self.dbmanager.PutRequest(request)
        return True
        
    def OnUpdateTaskTemplateRespond(self, respond):
        if respond.result != True:
            logging.warn('load task template %s fail' % respond.templateid)
            if respond.templateid in self.waitingruntemplate:
                self.waitingruntemplate.pop(respond.templateid)
            return 
        
        if self.LoadTaskTemplate(respond.info) == True:
            if respond.templateid in self.waitingruntemplate:
                info = self.waitingruntemplate.pop(respond.templateid)
                self.NewTask(respond.templateid, info['sql'])
                print 'scheduler, OnUpdateTaskTemplateRespond', datetime.datetime.now()
        else:
            logging.warn('load task template %s fail' % respond.templateid)
        
    
    def DeleteTaskTemplate(self, id):
        if id in self.task_templates:
            self.task_templates.pop(id)
            logging.info('delete tasktemplate %d' % id)
            return True
        return False
    
    
    def UpdateSubTask(self, subtasks):
        request = db.dbprotocol.UpdateSubTasksInfo()
        for v in subtasks:
            info = {}
            info['TaskStatusId'] = v.id
            info['TaskId'] = v.task.id
            info['Ip'] = v.ip
            info['MacAddress'] = v.mac_addr
            if v.start_time != None:
                info['StartTime'] = v.start_time
            
            if v.end_time != None:
                info['EndTime'] = v.end_time
            request.infos.append(info)
        self.dbmanager.PutRequest(request)
    
 
    def ExcuteTemplateTask(self, templateid, casesql=None, call_back=None):
        result = False
        if templateid not in self.task_templates:
            self.waitingruntemplate[templateid] = {'sql':casesql}
            self.UpdateTaskTemplate(templateid)
        else:
            result, new_task = self.NewTask(templateid, casesql, None)
            
        return result
    
    
    def TerminatedTask(self, taskid):
        if taskid not in self.tasks:
            logging.warn('cancle task not exits')
            return False
        
        t = self.tasks[taskid]
        if t.status != task.RUNING:
            return False
        machine_run_tasks = t.Terminat()
        for k, v in machine_run_tasks.items():
            c = self.clientmanager.GetClient(k)
            self.DoTerminatedTask(c.connectid, taskid, v)
        return True
    
    def OnRuningTaskInfoRequest(self, clientid, respond):
        for v in respond.tasklist:
            if v.tid not in self.tasks:
                logging.warn('taskid not exits')
                continue
        
            t = self.tasks[v.tid]
            
            if v.subid not in t.subtasks:
                logging.warn('subtaskid not exits')
                continue
    
            subtask = t.subtasks[v.subid]
            
            mechine = self.clientmanager.GetClient(clientid)
            
            
            if subtask.run_machine != mechine.dbid:
                logging.warn('machine id not match')
                continue
            
            mechine.AddTask(t.template, v.tid, v.subid)
            subtask.start_disconnect_time = None

    def GetTaskFromTempId(self, tid):
        for v in self.tasks.itervalues():
            if v.tempid == tid:
                return v
        return None